package com.shijunhao.graduation.utils;

import javafx.util.Pair;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.PriorityQueue;

/**
 * @author: shijunhao
 * @date: 2022/5/8 5:01
 * @description:
 */
public class TopKQueue {//自定义固定长度(k)的优先队列，因此可以解决Top-k问题
    PriorityQueue<Pair<Long,Double>> queue;
    int k;

    public TopKQueue(int k) {
        this.k = k;
        this.queue = new PriorityQueue<>(k, Comparator.comparingDouble(Pair::getValue));
    }

    public void add(Pair<Long,Double> e) {
        if (queue.size() < k) { //当前队列元素个数不足k个时，直接添加
            queue.add(e);
        } else { //超出k个时
            if (e.getValue().compareTo(queue.peek().getValue()) > 0) {// 如果新元素大于了堆顶元素，说明新元素应替换掉当前堆顶元素
                queue.poll();
                queue.add(e);
            }
        }
    }

    public List<Long> res(){
        List<Long> result = new ArrayList<>();
        while(!queue.isEmpty()){
            result.add(queue.poll().getKey());
        }
        return result;
    }
}
